import math
import time
import warnings

import numpy as np
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tsa.stattools import adfuller
import statsmodels.tsa.stattools as st

from app.utils import MathUtil


class ARIMAMethod:

    def __init__(self, fail_dataframe):
        self.fail_dataframe = fail_dataframe

    def get_window_matrix(self, length):
        segment = self.fail_dataframe["value"].tolist()
        t_matrix = []

        for item in segment:
            st = segment.index(item)
            if st > len(segment) - length:
                break

            row = []
            for i in range(st, st + length):
                row.append(segment[i])

            t_matrix.append(row)

        return np.array(t_matrix).T

    @staticmethod
    def get_SSA(mat):
        U, s, VT = np.linalg.svd(mat)
        con = s / np.sum(s)
        r = np.linalg.matrix_rank(mat)

        mats = []
        for i in range(r):
            item = s[i] * U[:, i][:, None] @ VT[i, :][:, None].T
            mats.append(item)

        seqs = []
        for mat in mats:
            mat = np.asarray(mat)
            row, col = mat.shape
            subSeq = [np.mean(mat[::-1, :].diagonal(i)) for i in range(1 - row, col)]
            seqs.append(subSeq)

        return seqs, con

    def get_data_set(self, data, folds):
        data = np.array(data)
        sp = math.floor(len(data) * folds)
        d_train = np.hsplit(data, (0, sp))[1]
        d_test = np.hsplit(data, [sp + 1, len(data)])[1]
        return d_train, d_test

    def model_arima(self, data, folds=0.8):
        warnings.filterwarnings("ignore")
        diff = 0

        mx = max(data)
        mn = min(data)
        dta = MathUtil.normalize(data, mx, mn)

        # adf & diff
        tmp = dta
        while adfuller(tmp)[1] > 0.05:
            diff += 1
            head = tmp[0]
            tmp = [tmp[i + 1] - tmp[i] for i in range(0, len(tmp) - 1)]
            tmp.insert(0, head)

        order_analyze = st.arma_order_select_ic(dta, max_ar=3, max_ma=3, ic='aic')

        # train and predict
        p = order_analyze['aic_min_order'][0]
        q = order_analyze['aic_min_order'][1]
        d_train, d_test = self.get_data_set(dta, folds)
        md = ARIMA(d_train, order=(p, diff, q)).fit()
        predict = md.forecast(len(d_test))

        # normalized
        predict = MathUtil.anti_normalize(predict, mx, mn)
        d_test = MathUtil.anti_normalize(d_test, mx, mn)

        return d_test, predict
